//! GridGraph2D4CMT - Multi-threaded 2D grid graph with 4-connectivity
//!
//! This module implements a parallel version of the 2D grid graph where each node
//! connects to its 4 orthogonal neighbors (left, right, up, down).
//!
//! The multi-threaded implementation uses:
//! - Rayon for parallel iteration and work distribution
//! - Atomic operations for thread-safe state updates
//! - Lock-free queues for efficient work sharing
//!
//! This variant is recommended for large grids (>1000x1000) where the overhead
//! of parallelization is offset by the performance gains from using multiple cores.

use crate::internal::memory::NodeIndex2D;
use crate::types::{GridCutError, Label, Parent, Segment};
use crossbeam::queue::SegQueue;
use num_traits::{Bounded, Num};
use rayon::prelude::*;
use std::sync::atomic::{AtomicU32, AtomicU8, Ordering};
use std::sync::Arc;

/// Multi-threaded 2D grid graph with 4-connected neighbors
///
/// This structure represents a regular 2D grid where each node can connect
/// to up to 4 neighbors (left, right, up, down). It implements a parallel
/// version of the Boykov-Kolmogorov maximum flow algorithm optimized for
/// multi-core processors.
///
/// # Type Parameters
///
/// * `TCap` - Terminal capacity type (edge capacities to source/sink)
/// * `NCap` - Neighbor capacity type (edge capacities between grid nodes)
/// * `Flow` - Flow accumulation type (must be able to hold sums of capacities)
///
/// # Example
///
/// ```ignore
/// use grid_cut_rs::GridGraph2D4CMT;
///
/// let mut graph = GridGraph2D4CMT::<i32, i32, i32>::new(100, 100)?;
/// let node = graph.node_id(50, 50)?;
/// graph.set_terminal_cap(node, 10, 5)?;
/// graph.compute_maxflow();
/// let segment = graph.get_segment(node);
/// ```
#[derive(Debug)]
pub struct GridGraph2D4CMT<TCap, NCap, Flow>
where
    TCap: Num + Copy + PartialOrd + Send + Sync,
    NCap: Num + Copy + PartialOrd + Send + Sync,
    Flow: Num + Copy + PartialOrd + From<TCap> + From<NCap> + Send + Sync,
{
    // Grid dimensions
    width: usize,
    height: usize,
    padded_width: usize,
    padded_height: usize,
    width_blocks: usize,

    // Node indexing helper
    indexer: NodeIndex2D,

    // Node state vectors (indexed by node_id) - using atomic types for thread safety
    labels: Vec<AtomicU8>,
    parents: Vec<AtomicU8>,
    parent_ids: Vec<AtomicU32>,
    timestamps: Vec<AtomicU32>,

    // Edge capacity vectors - protected by synchronization during updates
    rc: [Vec<NCap>; 4],
    rc_st: Vec<TCap>,

    // Thread-safe queues for parallel processing
    active_queue: Arc<SegQueue<usize>>,
    orphan_queue: Arc<SegQueue<usize>>,
    free_queue: Arc<SegQueue<usize>>,

    // Synchronization primitives
    num_threads: usize,

    // Algorithm state
    time: AtomicU32,
    max_flow: Flow,
}

// Helper functions for atomic label operations
impl Label {
    fn to_u8(self) -> u8 {
        match self {
            Label::Free => 0,
            Label::Source => 1,
            Label::Sink => 2,
        }
    }

    fn from_u8(val: u8) -> Self {
        match val {
            0 => Label::Free,
            1 => Label::Source,
            2 => Label::Sink,
            _ => Label::Free,
        }
    }
}

// Helper functions for atomic parent operations
impl Parent {
    fn to_u8(self) -> u8 {
        match self {
            Parent::GreaterX => 0,
            Parent::GreaterY => 1,
            Parent::LessY => 2,
            Parent::LessX => 3,
            Parent::None => 4,
            Parent::Terminal => 5,
        }
    }

    fn from_u8(val: u8) -> Self {
        match val {
            0 => Parent::GreaterX,
            1 => Parent::GreaterY,
            2 => Parent::LessY,
            3 => Parent::LessX,
            4 => Parent::None,
            5 => Parent::Terminal,
            _ => Parent::None,
        }
    }
}

impl<TCap, NCap, Flow> GridGraph2D4CMT<TCap, NCap, Flow>
where
    TCap: Num + Copy + PartialOrd + Bounded + std::ops::Neg<Output = TCap> + From<NCap> + Send + Sync,
    NCap: Num + Copy + PartialOrd + Bounded + From<TCap> + Send + Sync,
    Flow: Num + Copy + PartialOrd + From<TCap> + From<NCap> + Send + Sync,
{
    /// Creates a new multi-threaded 2D 4-connected grid graph
    ///
    /// # Arguments
    ///
    /// * `width` - Grid width (must be > 0)
    /// * `height` - Grid height (must be > 0)
    ///
    /// # Returns
    ///
    /// Returns `Ok(GridGraph2D4CMT)` on success, or `Err(GridCutError)` if:
    /// - Width or height is 0
    /// - Memory allocation fails
    pub fn new(width: usize, height: usize) -> Result<Self, GridCutError> {
        // Validate dimensions
        if width == 0 || height == 0 {
            return Err(GridCutError::InvalidDimensions { width, height });
        }

        // Calculate padded dimensions
        let padded_width = crate::internal::memory::next_higher_mul8(width + 2);
        let padded_height = crate::internal::memory::next_higher_mul8(height + 2);
        let width_blocks = padded_width >> 3;

        // Create node indexer
        let indexer = NodeIndex2D::new(padded_width, padded_height);

        // Calculate total number of nodes
        let num_nodes = padded_width * padded_height;

        // Initialize atomic node state vectors
        let labels = (0..num_nodes)
            .map(|_| AtomicU8::new(Label::Free.to_u8()))
            .collect();
        let parents = (0..num_nodes)
            .map(|_| AtomicU8::new(Parent::None.to_u8()))
            .collect();
        let parent_ids = (0..num_nodes).map(|_| AtomicU32::new(0)).collect();
        let timestamps = (0..num_nodes).map(|_| AtomicU32::new(0)).collect();

        // Initialize capacity vectors
        let rc = [
            vec![NCap::zero(); num_nodes],
            vec![NCap::zero(); num_nodes],
            vec![NCap::zero(); num_nodes],
            vec![NCap::zero(); num_nodes],
        ];
        let rc_st = vec![TCap::zero(); num_nodes];

        // Initialize thread-safe queues
        let active_queue = Arc::new(SegQueue::new());
        let orphan_queue = Arc::new(SegQueue::new());
        let free_queue = Arc::new(SegQueue::new());

        // Determine number of threads (use rayon's default)
        let num_threads = rayon::current_num_threads();

        Ok(Self {
            width,
            height,
            padded_width,
            padded_height,
            width_blocks,
            indexer,
            labels,
            parents,
            parent_ids,
            timestamps,
            rc,
            rc_st,
            active_queue,
            orphan_queue,
            free_queue,
            num_threads,
            time: AtomicU32::new(0),
            max_flow: Flow::zero(),
        })
    }

    /// Returns the node ID for grid coordinates (x, y)
    pub fn node_id(&self, x: usize, y: usize) -> Result<usize, GridCutError> {
        if x >= self.width || y >= self.height {
            return Err(GridCutError::InvalidCoordinates { x, y, z: None });
        }
        Ok(self.indexer.node_id(x + 1, y + 1))
    }

    /// Sets terminal capacities for a node
    pub fn set_terminal_cap(
        &mut self,
        node_id: usize,
        cap_source: TCap,
        cap_sink: TCap,
    ) -> Result<(), GridCutError> {
        if node_id >= self.labels.len() {
            return Err(GridCutError::InvalidNodeId {
                node_id,
                max: self.labels.len() - 1,
            });
        }

        // Set terminal capacity
        self.rc_st[node_id] = cap_source - cap_sink;

        // Initialize node based on terminal capacity
        if cap_source > cap_sink {
            self.labels[node_id].store(Label::Source.to_u8(), Ordering::Release);
            self.parents[node_id].store(Parent::Terminal.to_u8(), Ordering::Release);
            self.parent_ids[node_id].store(node_id as u32, Ordering::Release);
            self.free_queue.push(node_id);
        } else if cap_sink > cap_source {
            self.labels[node_id].store(Label::Sink.to_u8(), Ordering::Release);
            self.parents[node_id].store(Parent::Terminal.to_u8(), Ordering::Release);
            self.parent_ids[node_id].store(node_id as u32, Ordering::Release);
            self.free_queue.push(node_id);
        }

        Ok(())
    }

    /// Sets neighbor edge capacity
    pub fn set_neighbor_cap(
        &mut self,
        node_id: usize,
        offset_x: i32,
        offset_y: i32,
        cap: NCap,
    ) -> Result<(), GridCutError> {
        if node_id >= self.labels.len() {
            return Err(GridCutError::InvalidNodeId {
                node_id,
                max: self.labels.len() - 1,
            });
        }

        let direction = match (offset_x, offset_y) {
            (1, 0) => 0,
            (0, 1) => 1,
            (0, -1) => 2,
            (-1, 0) => 3,
            _ => {
                return Err(GridCutError::InvalidOffset {
                    offset_x,
                    offset_y,
                    offset_z: None,
                });
            }
        };

        self.rc[direction][node_id] = cap;
        Ok(())
    }

    /// Bulk set all capacities from arrays
    pub fn set_caps(
        &mut self,
        cap_source: &[TCap],
        cap_sink: &[TCap],
        cap_le: &[NCap],
        cap_ge: &[NCap],
        cap_el: &[NCap],
        cap_eg: &[NCap],
    ) -> Result<(), GridCutError> {
        let expected_len = self.width * self.height;

        for (_name, len) in [
            ("cap_source", cap_source.len()),
            ("cap_sink", cap_sink.len()),
            ("cap_le", cap_le.len()),
            ("cap_ge", cap_ge.len()),
            ("cap_el", cap_el.len()),
            ("cap_eg", cap_eg.len()),
        ] {
            if len != expected_len {
                return Err(GridCutError::InvalidCapacityArrayLength {
                    expected: expected_len,
                    actual: len,
                });
            }
        }

        for y in 0..self.height {
            for x in 0..self.width {
                let idx = y * self.width + x;
                let node = self.indexer.node_id(x + 1, y + 1);

                self.rc_st[node] = cap_source[idx] - cap_sink[idx];

                if cap_source[idx] > cap_sink[idx] {
                    self.labels[node].store(Label::Source.to_u8(), Ordering::Release);
                    self.parents[node].store(Parent::Terminal.to_u8(), Ordering::Release);
                    self.parent_ids[node].store(node as u32, Ordering::Release);
                    self.free_queue.push(node);
                } else if cap_sink[idx] > cap_source[idx] {
                    self.labels[node].store(Label::Sink.to_u8(), Ordering::Release);
                    self.parents[node].store(Parent::Terminal.to_u8(), Ordering::Release);
                    self.parent_ids[node].store(node as u32, Ordering::Release);
                    self.free_queue.push(node);
                }

                self.rc[0][node] = cap_ge[idx];
                self.rc[1][node] = cap_eg[idx];
                self.rc[2][node] = cap_el[idx];
                self.rc[3][node] = cap_le[idx];
            }
        }

        Ok(())
    }

    /// Returns the segment (Source or Sink) that a node belongs to
    pub fn get_segment(&self, node_id: usize) -> Segment {
        let label = Label::from_u8(self.labels[node_id].load(Ordering::Relaxed));
        match label {
            Label::Source => Segment::Source,
            Label::Sink => Segment::Sink,
            Label::Free => Segment::Sink,
        }
    }

    /// Returns the computed maximum flow value
    pub fn get_flow(&self) -> Flow {
        self.max_flow
    }

    // Helper methods
    #[inline]
    fn get_neighbor(&self, node: usize, direction: usize) -> usize {
        match direction {
            0 => self.indexer.neighbor_ge(node),
            1 => self.indexer.neighbor_eg(node),
            2 => self.indexer.neighbor_el(node),
            3 => self.indexer.neighbor_le(node),
            _ => unreachable!("Invalid direction"),
        }
    }

    #[inline]
    fn sister(direction: usize) -> usize {
        match direction {
            0 => 3,
            1 => 2,
            2 => 1,
            3 => 0,
            _ => unreachable!("Invalid direction"),
        }
    }

    #[inline]
    fn direction_to_parent(direction: usize) -> Parent {
        match direction {
            0 => Parent::GreaterX,
            1 => Parent::GreaterY,
            2 => Parent::LessY,
            3 => Parent::LessX,
            _ => unreachable!("Invalid direction"),
        }
    }
}


// Parallel algorithm implementation
impl<TCap, NCap, Flow> GridGraph2D4CMT<TCap, NCap, Flow>
where
    TCap: Num + Copy + PartialOrd + Bounded + std::ops::Neg<Output = TCap> + From<NCap> + Send + Sync,
    NCap: Num + Copy + PartialOrd + Bounded + From<TCap> + Send + Sync,
    Flow: Num + Copy + PartialOrd + From<TCap> + From<NCap> + Send + Sync,
{
    /// Parallel growth phase: searches for an augmenting path from source to sink
    ///
    /// Multiple threads work concurrently to expand the search trees.
    /// Returns Some((source_node, sink_node, connecting_direction)) if a path is found.
    fn parallel_grow(&self) -> Option<(usize, usize, usize)> {
        use std::sync::atomic::AtomicBool;
        
        // Shared flag to signal when a path is found
        let path_found = Arc::new(AtomicBool::new(false));
        let result = Arc::new(SegQueue::new());

        // Process nodes in parallel
        let active_queue = Arc::clone(&self.active_queue);
        let path_found_clone = Arc::clone(&path_found);
        let result_clone = Arc::clone(&result);

        // Use rayon to process nodes in parallel
        (0..self.num_threads).into_par_iter().for_each(|_| {
            while !path_found_clone.load(Ordering::Relaxed) {
                // Try to get a node from the active queue
                let node = match active_queue.pop() {
                    Some(n) => n,
                    None => break, // No more nodes to process
                };

                let label = Label::from_u8(self.labels[node].load(Ordering::Acquire));

                // Try all 4 neighbor directions
                for direction in 0..4 {
                    if path_found_clone.load(Ordering::Relaxed) {
                        break;
                    }

                    let neighbor = self.get_neighbor(node, direction);

                    if label == Label::Source {
                        // Expanding source tree
                        if self.rc[direction][node] > NCap::zero() {
                            let neighbor_label =
                                Label::from_u8(self.labels[neighbor].load(Ordering::Acquire));

                            match neighbor_label {
                                Label::Sink => {
                                    // Found augmenting path!
                                    if !path_found_clone.swap(true, Ordering::SeqCst) {
                                        result_clone.push((node, neighbor, direction));
                                    }
                                    break;
                                }
                                Label::Free => {
                                    // Try to claim this node for source tree
                                    if self.labels[neighbor]
                                        .compare_exchange(
                                            Label::Free.to_u8(),
                                            Label::Source.to_u8(),
                                            Ordering::SeqCst,
                                            Ordering::Relaxed,
                                        )
                                        .is_ok()
                                    {
                                        // Successfully claimed - set parent and add to queue
                                        self.parents[neighbor].store(
                                            Self::direction_to_parent(Self::sister(direction))
                                                .to_u8(),
                                            Ordering::Release,
                                        );
                                        self.parent_ids[neighbor]
                                            .store(node as u32, Ordering::Release);
                                        active_queue.push(neighbor);
                                    }
                                }
                                Label::Source => {
                                    // Already in source tree
                                }
                            }
                        }
                    } else if label == Label::Sink {
                        // Expanding sink tree
                        let sister_dir = Self::sister(direction);
                        if self.rc[sister_dir][neighbor] > NCap::zero() {
                            let neighbor_label =
                                Label::from_u8(self.labels[neighbor].load(Ordering::Acquire));

                            match neighbor_label {
                                Label::Source => {
                                    // Found augmenting path!
                                    if !path_found_clone.swap(true, Ordering::SeqCst) {
                                        result_clone.push((neighbor, node, sister_dir));
                                    }
                                    break;
                                }
                                Label::Free => {
                                    // Try to claim this node for sink tree
                                    if self.labels[neighbor]
                                        .compare_exchange(
                                            Label::Free.to_u8(),
                                            Label::Sink.to_u8(),
                                            Ordering::SeqCst,
                                            Ordering::Relaxed,
                                        )
                                        .is_ok()
                                    {
                                        // Successfully claimed - set parent and add to queue
                                        self.parents[neighbor].store(
                                            Self::direction_to_parent(Self::sister(direction))
                                                .to_u8(),
                                            Ordering::Release,
                                        );
                                        self.parent_ids[neighbor]
                                            .store(node as u32, Ordering::Release);
                                        active_queue.push(neighbor);
                                    }
                                }
                                Label::Sink => {
                                    // Already in sink tree
                                }
                            }
                        }
                    }
                }
            }
        });

        // Return the first path found (if any)
        result.pop()
    }
}


impl<TCap, NCap, Flow> GridGraph2D4CMT<TCap, NCap, Flow>
where
    TCap: Num + Copy + PartialOrd + Bounded + std::ops::Neg<Output = TCap> + From<NCap> + Send + Sync,
    NCap: Num + Copy + PartialOrd + Bounded + From<TCap> + Send + Sync,
    Flow: Num + Copy + PartialOrd + From<TCap> + From<NCap> + Send + Sync,
{
    /// Find minimum residual capacity along source tree path
    fn find_minrf_s(&self, mut v: usize) -> NCap {
        let mut min_cap = NCap::max_value();

        while Parent::from_u8(self.parents[v].load(Ordering::Acquire)) != Parent::Terminal {
            let parent = Parent::from_u8(self.parents[v].load(Ordering::Acquire));
            let parent_id = self.parent_ids[v].load(Ordering::Acquire) as usize;

            let direction = match parent {
                Parent::GreaterX => 0,
                Parent::GreaterY => 1,
                Parent::LessY => 2,
                Parent::LessX => 3,
                _ => break,
            };

            let sister_dir = Self::sister(direction);
            let cap = self.rc[sister_dir][parent_id];

            if cap < min_cap {
                min_cap = cap;
            }

            v = parent_id;
        }

        // Check terminal capacity
        let terminal_cap = self.rc_st[v];
        if terminal_cap > TCap::zero() {
            let terminal_cap_n = NCap::from(terminal_cap);
            if terminal_cap_n < min_cap {
                min_cap = terminal_cap_n;
            }
        }

        min_cap
    }

    /// Find minimum residual capacity along sink tree path
    fn find_minrf_t(&self, mut v: usize) -> NCap {
        let mut min_cap = NCap::max_value();

        while Parent::from_u8(self.parents[v].load(Ordering::Acquire)) != Parent::Terminal {
            let parent = Parent::from_u8(self.parents[v].load(Ordering::Acquire));
            let parent_id = self.parent_ids[v].load(Ordering::Acquire) as usize;

            let direction = match parent {
                Parent::GreaterX => 0,
                Parent::GreaterY => 1,
                Parent::LessY => 2,
                Parent::LessX => 3,
                _ => break,
            };

            let sister_dir = Self::sister(direction);
            let cap = self.rc[sister_dir][parent_id];

            if cap < min_cap {
                min_cap = cap;
            }

            v = parent_id;
        }

        // Check terminal capacity
        let terminal_cap = self.rc_st[v];
        if terminal_cap < TCap::zero() {
            let terminal_cap_n = NCap::from(-terminal_cap);
            if terminal_cap_n < min_cap {
                min_cap = terminal_cap_n;
            }
        }

        min_cap
    }

    /// Augment flow along source tree path (sequential - called with exclusive access)
    fn aug_s(&mut self, mut v: usize, bottleneck: NCap) {
        while Parent::from_u8(self.parents[v].load(Ordering::Acquire)) != Parent::Terminal {
            let parent = Parent::from_u8(self.parents[v].load(Ordering::Acquire));
            let parent_id = self.parent_ids[v].load(Ordering::Acquire) as usize;

            let direction = match parent {
                Parent::GreaterX => 0,
                Parent::GreaterY => 1,
                Parent::LessY => 2,
                Parent::LessX => 3,
                _ => break,
            };

            let sister_dir = Self::sister(direction);

            // Update residual capacities
            self.rc[sister_dir][parent_id] = self.rc[sister_dir][parent_id] - bottleneck;
            self.rc[direction][v] = self.rc[direction][v] + bottleneck;

            // If edge saturated, mark node as orphan
            if self.rc[sister_dir][parent_id] == NCap::zero() {
                self.parents[v].store(Parent::None.to_u8(), Ordering::Release);
                self.orphan_queue.push(v);
            }

            v = parent_id;
        }

        // Update terminal capacity
        let bottleneck_t = TCap::from(bottleneck);
        self.rc_st[v] = self.rc_st[v] - bottleneck_t;

        // If terminal edge saturated, mark node as orphan
        if self.rc_st[v] == TCap::zero() {
            self.parents[v].store(Parent::None.to_u8(), Ordering::Release);
            self.orphan_queue.push(v);
        }
    }

    /// Augment flow along sink tree path (sequential - called with exclusive access)
    fn aug_t(&mut self, mut v: usize, bottleneck: NCap) {
        while Parent::from_u8(self.parents[v].load(Ordering::Acquire)) != Parent::Terminal {
            let parent = Parent::from_u8(self.parents[v].load(Ordering::Acquire));
            let parent_id = self.parent_ids[v].load(Ordering::Acquire) as usize;

            let direction = match parent {
                Parent::GreaterX => 0,
                Parent::GreaterY => 1,
                Parent::LessY => 2,
                Parent::LessX => 3,
                _ => break,
            };

            let sister_dir = Self::sister(direction);

            // Update residual capacities
            self.rc[sister_dir][parent_id] = self.rc[sister_dir][parent_id] - bottleneck;
            self.rc[direction][v] = self.rc[direction][v] + bottleneck;

            // If edge saturated, mark node as orphan
            if self.rc[sister_dir][parent_id] == NCap::zero() {
                self.parents[v].store(Parent::None.to_u8(), Ordering::Release);
                self.orphan_queue.push(v);
            }

            v = parent_id;
        }

        // Update terminal capacity
        let bottleneck_t = TCap::from(bottleneck);
        self.rc_st[v] = self.rc_st[v] + bottleneck_t;

        // If terminal edge saturated, mark node as orphan
        if self.rc_st[v] == TCap::zero() {
            self.parents[v].store(Parent::None.to_u8(), Ordering::Release);
            self.orphan_queue.push(v);
        }
    }

    /// Augmentation phase: push flow along the augmenting path
    /// This is done sequentially as it modifies shared capacity arrays
    fn augment(&mut self, vs: usize, vt: usize, st: usize) {
        // Find minimum residual capacity (bottleneck)
        let mut bottleneck = self.rc[st][vs];

        let min_s = self.find_minrf_s(vs);
        if min_s < bottleneck {
            bottleneck = min_s;
        }

        let min_t = self.find_minrf_t(vt);
        if min_t < bottleneck {
            bottleneck = min_t;
        }

        // Update edge connecting the two trees
        self.rc[st][vs] = self.rc[st][vs] - bottleneck;
        let sister_st = Self::sister(st);
        self.rc[sister_st][vt] = self.rc[sister_st][vt] + bottleneck;

        // Augment along source tree path
        self.aug_s(vs, bottleneck);

        // Augment along sink tree path
        self.aug_t(vt, bottleneck);

        // Accumulate flow
        self.max_flow = self.max_flow + Flow::from(bottleneck);
    }

    /// Find origin (terminal) of a tree path using timestamps
    fn find_origin(&self, start: usize) -> Parent {
        let mut v = start;
        let current_time = self.time.load(Ordering::Acquire);

        // Follow parent pointers until we reach a terminal or detect invalid path
        let mut visited: Vec<usize> = Vec::new();
        while Parent::from_u8(self.parents[v].load(Ordering::Acquire)) != Parent::Terminal
            && Parent::from_u8(self.parents[v].load(Ordering::Acquire)) != Parent::None
        {
            // Check if this node was already validated in current timestamp
            if self.timestamps[v].load(Ordering::Acquire) == current_time {
                // Mark all nodes on this path
                for node_idx in visited.iter() {
                    self.timestamps[*node_idx].store(current_time, Ordering::Release);
                }
                return Parent::Terminal;
            }

            visited.push(v);
            v = self.parent_ids[v].load(Ordering::Acquire) as usize;

            // Safety check to prevent infinite loops
            if visited.len() > self.labels.len() {
                return Parent::None;
            }
        }

        // Mark all nodes on this path with current timestamp if valid
        let result = Parent::from_u8(self.parents[v].load(Ordering::Acquire));
        if result == Parent::Terminal {
            for node_idx in visited.iter() {
                self.timestamps[*node_idx].store(current_time, Ordering::Release);
            }
            self.timestamps[v].store(current_time, Ordering::Release);
        }

        result
    }

    /// Adoption phase: restore tree structure after augmentation
    /// This is done sequentially to maintain consistency
    fn adopt(&mut self) {
        let orphan_queue2 = SegQueue::new();

        while !self.orphan_queue.is_empty() || !orphan_queue2.is_empty() {
            let v = if !orphan_queue2.is_empty() {
                match orphan_queue2.pop() {
                    Some(n) => n,
                    None => continue,
                }
            } else {
                match self.orphan_queue.pop() {
                    Some(n) => n,
                    None => break,
                }
            };

            let label = Label::from_u8(self.labels[v].load(Ordering::Acquire));

            // Try to find a new valid parent
            let mut found_parent = false;

            // First, check if this node can reconnect directly to terminal
            let has_terminal_capacity = if label == Label::Source {
                self.rc_st[v] > TCap::zero()
            } else {
                self.rc_st[v] < TCap::zero()
            };

            if has_terminal_capacity {
                // Can reconnect to terminal
                self.parents[v].store(Parent::Terminal.to_u8(), Ordering::Release);
                self.parent_ids[v].store(v as u32, Ordering::Release);
                self.timestamps[v].store(self.time.load(Ordering::Acquire), Ordering::Release);
                found_parent = true;
            } else {
                // Try to find a neighbor parent
                for direction in 0..4 {
                    let neighbor = self.get_neighbor(v, direction);

                    // Check if neighbor can be a valid parent
                    if Label::from_u8(self.labels[neighbor].load(Ordering::Acquire)) == label {
                        let sister_dir = Self::sister(direction);

                        // Check if edge from neighbor to v has capacity
                        if self.rc[sister_dir][neighbor] > NCap::zero() {
                            // Check if neighbor has valid path to terminal
                            if self.find_origin(neighbor) == Parent::Terminal {
                                // Found valid parent!
                                self.parents[v].store(
                                    Self::direction_to_parent(direction).to_u8(),
                                    Ordering::Release,
                                );
                                self.parent_ids[v].store(neighbor as u32, Ordering::Release);
                                self.timestamps[v].store(
                                    self.time.load(Ordering::Acquire),
                                    Ordering::Release,
                                );
                                found_parent = true;
                                break;
                            }
                        }
                    }
                }
            }

            if !found_parent {
                // Could not find valid parent - make node free
                self.labels[v].store(Label::Free.to_u8(), Ordering::Release);
                self.free_queue.push(v);

                // Add neighbors that had v as parent to orphan queue
                for direction in 0..4 {
                    let neighbor = self.get_neighbor(v, direction);

                    if Label::from_u8(self.labels[neighbor].load(Ordering::Acquire)) == label {
                        let sister_dir = Self::sister(direction);
                        let expected_parent = Self::direction_to_parent(sister_dir);

                        if Parent::from_u8(self.parents[neighbor].load(Ordering::Acquire))
                            == expected_parent
                            && self.parent_ids[neighbor].load(Ordering::Acquire) == v as u32
                        {
                            self.parents[neighbor].store(Parent::None.to_u8(), Ordering::Release);
                            orphan_queue2.push(neighbor);
                        }
                    }
                }
            }
        }

        // Re-activate neighbors of freed nodes
        while let Some(v) = self.free_queue.pop() {
            for direction in 0..4 {
                let neighbor = self.get_neighbor(v, direction);
                let label = Label::from_u8(self.labels[neighbor].load(Ordering::Acquire));

                // Only reactivate nodes that are in a tree
                if (label == Label::Source || label == Label::Sink)
                    && Parent::from_u8(self.parents[neighbor].load(Ordering::Acquire))
                        != Parent::None
                {
                    // Check if this neighbor can now expand into the freed node v
                    let can_expand = if label == Label::Source {
                        self.rc[direction][neighbor] > NCap::zero()
                    } else {
                        let sister_dir = Self::sister(direction);
                        self.rc[sister_dir][neighbor] > NCap::zero()
                    };

                    if can_expand {
                        self.active_queue.push(neighbor);
                    }
                }
            }
        }
    }

    /// Computes the maximum flow from source to sink using parallel algorithm
    pub fn compute_maxflow(&mut self) {
        // Initialize active queue from free_queue
        while let Some(v) = self.free_queue.pop() {
            self.active_queue.push(v);
        }

        // Main algorithm loop
        loop {
            // Parallel growth phase: find augmenting path
            let path = match self.parallel_grow() {
                Some(p) => p,
                None => break, // No more augmenting paths
            };

            // Increment timestamp
            self.time.fetch_add(1, Ordering::SeqCst);

            // Sequential augmentation phase: push flow along path
            let (vs, vt, st) = path;
            self.augment(vs, vt, st);

            // Sequential adoption phase: restore tree structure
            self.adopt();
            
            // Memory fence to ensure all updates are visible before next growth phase
            std::sync::atomic::fence(Ordering::SeqCst);
        }

        // Final segmentation pass
        self.finalize_segmentation();
    }

    /// Final segmentation pass after maxflow completes
    fn finalize_segmentation(&mut self) {
        // Convert Free nodes to Sink
        for label in &self.labels {
            let current = label.load(Ordering::Relaxed);
            if current == Label::Free.to_u8() {
                label.store(Label::Sink.to_u8(), Ordering::Relaxed);
            }
        }
    }
}

// Implement GridGraph trait for GridGraph2D4CMT
impl<TCap, NCap, Flow> crate::types::GridGraph for GridGraph2D4CMT<TCap, NCap, Flow>
where
    TCap: Num + Copy + PartialOrd + Bounded + std::ops::Neg<Output = TCap> + From<NCap> + Send + Sync,
    NCap: Num + Copy + PartialOrd + Bounded + From<TCap> + Send + Sync,
    Flow: Num + Copy + PartialOrd + From<TCap> + From<NCap> + Send + Sync,
{
    type TerminalCap = TCap;
    type NeighborCap = NCap;
    type Flow = Flow;

    fn compute_maxflow(&mut self) {
        self.compute_maxflow()
    }

    fn get_flow(&self) -> Self::Flow {
        self.get_flow()
    }

    fn get_segment(&self, node_id: usize) -> Segment {
        self.get_segment(node_id)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mt_constructor_valid_dimensions() {
        let graph = GridGraph2D4CMT::<i32, i32, i32>::new(10, 20);
        assert!(graph.is_ok());

        let graph = graph.unwrap();
        assert_eq!(graph.width, 10);
        assert_eq!(graph.height, 20);
        assert_eq!(graph.max_flow, 0);
    }

    #[test]
    fn test_mt_simple_2x2_grid() {
        // Test that multi-threaded version produces correct results
        let mut graph = GridGraph2D4CMT::<i32, i32, i32>::new(2, 2).unwrap();

        let n00 = graph.node_id(0, 0).unwrap();
        let n10 = graph.node_id(1, 0).unwrap();
        let n01 = graph.node_id(0, 1).unwrap();
        let n11 = graph.node_id(1, 1).unwrap();

        // Set terminal capacities
        graph.set_terminal_cap(n00, 10, 0).unwrap();
        graph.set_terminal_cap(n11, 0, 10).unwrap();

        // Set neighbor capacities
        graph.set_neighbor_cap(n00, 1, 0, 5).unwrap();
        graph.set_neighbor_cap(n00, 0, 1, 5).unwrap();
        graph.set_neighbor_cap(n10, 0, 1, 5).unwrap();
        graph.set_neighbor_cap(n01, 1, 0, 5).unwrap();

        // Compute maxflow
        graph.compute_maxflow();

        // Verify flow value
        let flow = graph.get_flow();
        assert_eq!(flow, 10, "Flow should be exactly 10");

        // Verify segmentation
        let seg_n00 = graph.get_segment(n00);
        let seg_n11 = graph.get_segment(n11);
        assert_ne!(seg_n00, seg_n11, "Source and sink should be in different segments");
    }

    #[test]
    fn test_mt_matches_single_threaded_simple() {
        // Compare multi-threaded results with single-threaded version
        use crate::grid_cut::grid_2d_4c::GridGraph2D4C;

        let width = 4;
        let height = 4;

        // Set up single-threaded version
        let mut graph_st = GridGraph2D4C::<i32, i32, i32>::new(width, height).unwrap();
        
        // Set up multi-threaded version
        let mut graph_mt = GridGraph2D4CMT::<i32, i32, i32>::new(width, height).unwrap();

        // Configure both graphs identically
        let source = graph_st.node_id(0, 0).unwrap();
        let sink = graph_st.node_id(3, 3).unwrap();

        graph_st.set_terminal_cap(source, 100, 0).unwrap();
        graph_st.set_terminal_cap(sink, 0, 100).unwrap();
        
        graph_mt.set_terminal_cap(source, 100, 0).unwrap();
        graph_mt.set_terminal_cap(sink, 0, 100).unwrap();

        // Set uniform neighbor capacities
        for y in 0..height {
            for x in 0..width {
                let node = graph_st.node_id(x, y).unwrap();

                if x < width - 1 {
                    graph_st.set_neighbor_cap(node, 1, 0, 10).unwrap();
                    graph_mt.set_neighbor_cap(node, 1, 0, 10).unwrap();
                }
                if y < height - 1 {
                    graph_st.set_neighbor_cap(node, 0, 1, 10).unwrap();
                    graph_mt.set_neighbor_cap(node, 0, 1, 10).unwrap();
                }
            }
        }

        // Compute maxflow on both
        graph_st.compute_maxflow();
        graph_mt.compute_maxflow();

        // Verify flows match
        assert_eq!(
            graph_st.get_flow(),
            graph_mt.get_flow(),
            "Multi-threaded and single-threaded flows should match"
        );

        // Verify segmentations match
        for y in 0..height {
            for x in 0..width {
                let node = graph_st.node_id(x, y).unwrap();
                assert_eq!(
                    graph_st.get_segment(node),
                    graph_mt.get_segment(node),
                    "Segmentation should match at ({}, {})",
                    x,
                    y
                );
            }
        }
    }

    #[test]
    fn test_mt_matches_single_threaded_complex() {
        // Test with a more complex graph
        use crate::grid_cut::grid_2d_4c::GridGraph2D4C;

        let width = 8;
        let height = 8;

        let mut graph_st = GridGraph2D4C::<i32, i32, i32>::new(width, height).unwrap();
        let mut graph_mt = GridGraph2D4CMT::<i32, i32, i32>::new(width, height).unwrap();

        // Create a checkerboard pattern of sources and sinks
        for y in 0..height {
            for x in 0..width {
                let node = graph_st.node_id(x, y).unwrap();
                
                if (x + y) % 2 == 0 {
                    graph_st.set_terminal_cap(node, 50, 0).unwrap();
                    graph_mt.set_terminal_cap(node, 50, 0).unwrap();
                } else {
                    graph_st.set_terminal_cap(node, 0, 50).unwrap();
                    graph_mt.set_terminal_cap(node, 0, 50).unwrap();
                }

                // Set neighbor capacities
                if x < width - 1 {
                    graph_st.set_neighbor_cap(node, 1, 0, 15).unwrap();
                    graph_mt.set_neighbor_cap(node, 1, 0, 15).unwrap();
                }
                if y < height - 1 {
                    graph_st.set_neighbor_cap(node, 0, 1, 15).unwrap();
                    graph_mt.set_neighbor_cap(node, 0, 1, 15).unwrap();
                }
            }
        }

        graph_st.compute_maxflow();
        graph_mt.compute_maxflow();

        // Verify flows match
        assert_eq!(
            graph_st.get_flow(),
            graph_mt.get_flow(),
            "Multi-threaded and single-threaded flows should match for complex graph"
        );
    }

    #[test]
    fn test_mt_set_caps_bulk() {
        // Test bulk capacity setting
        let width = 3;
        let height = 2;
        let mut graph = GridGraph2D4CMT::<i32, i32, i32>::new(width, height).unwrap();

        let size = width * height;
        let cap_source = vec![100; size];
        let cap_sink = vec![50; size];
        let cap_le = vec![10; size];
        let cap_ge = vec![20; size];
        let cap_el = vec![30; size];
        let cap_eg = vec![40; size];

        let result = graph.set_caps(&cap_source, &cap_sink, &cap_le, &cap_ge, &cap_el, &cap_eg);
        assert!(result.is_ok());

        // Compute maxflow
        graph.compute_maxflow();

        // Should produce some flow
        assert!(graph.get_flow() > 0);
    }

    #[test]
    fn test_mt_no_flow_disconnected() {
        // Test a grid where source and sink are disconnected
        let mut graph = GridGraph2D4CMT::<i32, i32, i32>::new(3, 3).unwrap();

        let source = graph.node_id(0, 0).unwrap();
        let sink = graph.node_id(2, 2).unwrap();

        graph.set_terminal_cap(source, 100, 0).unwrap();
        graph.set_terminal_cap(sink, 0, 100).unwrap();

        // Don't set any neighbor capacities

        graph.compute_maxflow();

        // Flow should be zero
        assert_eq!(graph.get_flow(), 0);
    }

    #[test]
    fn test_mt_single_path() {
        // Test a grid with a single path
        let mut graph = GridGraph2D4CMT::<i32, i32, i32>::new(4, 1).unwrap();

        let n0 = graph.node_id(0, 0).unwrap();
        let n1 = graph.node_id(1, 0).unwrap();
        let n2 = graph.node_id(2, 0).unwrap();
        let n3 = graph.node_id(3, 0).unwrap();

        graph.set_terminal_cap(n0, 50, 0).unwrap();
        graph.set_terminal_cap(n3, 0, 50).unwrap();

        graph.set_neighbor_cap(n0, 1, 0, 30).unwrap();
        graph.set_neighbor_cap(n1, 1, 0, 20).unwrap();
        graph.set_neighbor_cap(n2, 1, 0, 40).unwrap();

        graph.compute_maxflow();

        // Flow should be limited by minimum capacity (20)
        assert_eq!(graph.get_flow(), 20);
    }

    #[test]
    fn test_mt_thread_count() {
        // Verify that the graph uses multiple threads
        let graph = GridGraph2D4CMT::<i32, i32, i32>::new(10, 10).unwrap();
        
        // Should use at least 1 thread (typically more on multi-core systems)
        assert!(graph.num_threads >= 1);
        
        println!("Multi-threaded graph using {} threads", graph.num_threads);
    }
}
